package com.dmai.capacity.platform.log.trace;

import java.util.Collections;
import java.util.concurrent.atomic.AtomicInteger;

import cn.hutool.core.lang.Snowflake;
import cn.hutool.extra.spring.SpringUtil;
import org.slf4j.MDC;

import com.alibaba.ttl.TransmittableThreadLocal;
import com.google.common.base.Optional;

import cn.hutool.core.util.IdUtil;
import cn.hutool.core.util.StrUtil;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.scripting.support.ResourceScriptSource;

/**
 * 日志追踪工具类
 */
public class MDCTraceUtils {
    /**
     * 追踪id的名称
     */
    public static final String KEY_TRACE_ID = "traceId";
    /**
     * 块id的名称
     */
    public static final String KEY_SPAN_ID = "spanId";

    /**
     * 日志链路追踪id信息头
     */
    public static final String TRACE_ID_HEADER = "x-traceId-header";
    /**
     * 日志链路块id信息头
     */
    public static final String SPAN_ID_HEADER = "x-spanId-header";

    /**
     * filter的优先级，值越低越优先
     */
    public static final int FILTER_ORDER = -1;

    private static final TransmittableThreadLocal<AtomicInteger> spanNumber = new TransmittableThreadLocal<>();

    private static StringRedisTemplate stringRedisTemplate = SpringUtil.getBean(StringRedisTemplate.class);
    private static String SNOWFLAKE_WORKID = "snowflake:workid";
    private static final SnowflakeIdWorker snowflake = new SnowflakeIdWorker(getWorkerId(SNOWFLAKE_WORKID));


    /**
     * 容器环境生成workid 并redis缓存
     *
     * @param key
     * @return
     */
    public static Long getWorkerId(String key) {
        DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>();
        redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("redis/redis_worker_id.lua")));
        redisScript.setResultType(Long.class);
        return stringRedisTemplate.execute(redisScript, Collections.singletonList(key));
    }

    /**
     * 创建traceId并赋值MDC
     */
    public static void addTrace() {
        String traceId = createTraceId();
        MDC.put(KEY_TRACE_ID, traceId);
        MDC.put(KEY_SPAN_ID, "0");
        initSpanNumber();
    }

    /**
     * 赋值MDC
     */
    public static void putTrace(String traceId, String spanId) {
        MDC.put(KEY_TRACE_ID, traceId);
        MDC.put(KEY_SPAN_ID, spanId);
        initSpanNumber();
    }

    /**
     * 获取MDC中的traceId值
     */
    public static String getTraceId() {
        return MDC.get(KEY_TRACE_ID);
    }

    /**
     * 获取MDC中的spanId值
     */
    public static String getSpanId() {
        return MDC.get(KEY_SPAN_ID);
    }

    /**
     * 清除MDC的值
     */
    public static void removeTrace() {
        MDC.remove(KEY_TRACE_ID);
        MDC.remove(KEY_SPAN_ID);
        spanNumber.remove();
    }

    /**
     * MongoDB唯一生成算法生成TraceId
     *
     * @return
     */
    public static String createTraceId() {
//		return  IdUtil.objectId() ;
        return String.valueOf(snowflake.nextId());
    }


    public static String getNextSpanId() {
        spanNumber.set(Optional.of(spanNumber.get()).or(new AtomicInteger(0)));
        return StrUtil.format("{}.{}", getSpanId(), spanNumber.get().incrementAndGet());
    }

    private static void initSpanNumber() {
        spanNumber.set(new AtomicInteger(0));
    }


}
